from utils.basic_function import *
from shapely.ops import snap, nearest_points


class SnapTool(PreprocessingTool):
    def __init__(self, table, tolerance, extend_threshold, extend_ratio, transform_distance, end_priority=True):
        super(SnapTool, self).__init__(table)
        self.tolerance = tolerance
        self.extend_threshold = extend_threshold
        self.extend_ratio = extend_ratio  # 次近路段与最近路段距离之比
        self.end_priority = end_priority
        self.transform_distance = transform_distance

    def snap_nodes_table(self):
        pk_list = self.table.get_pk_list()
        self.table.update_length()  # 更新length属性
        self.set_pk_list(pk_list)
        self.snap_nodes_workflow()

    def snap_nodes_workflow(self):
        count = 0
        for pk in self.pk_list:
            print(pk)
            geo_series = self.table.select_by_attr(self.table.pk, pk).iloc[0]
            geom = geo_series[self.table.geom]
            st_coord = geom.coords[0]
            ed_coord = geom.coords[-1]

            modified = self.snap_node(pk, st_coord, geom)

            if modified:  # 获取更新后的geom
                geo_series = self.table.select_by_attr(self.table.pk, pk).iloc[0]
                geom = geo_series[self.table.geom]
            self.snap_node(pk, ed_coord, geom)

            count += 1
            if count == 500:
                self.table.vacuum_analyze()
                count = 0

    def snap_node(self, pk_self, coord, geom_self):
        pk_col = self.table.pk
        geom_col = self.table.geom
        gdf = self.table.buffer_nodes(Point(coord), self.tolerance, and_where=f"and {pk_col} != {pk_self}")
        modified = False
        if not gdf.empty:
            # 附近有结点
            with conn:
                for i in range(gdf.shape[0]):
                    pt = Point(coord)
                    geom_neighbor = gdf[geom_col].iloc[i]
                    if pt.equals(Point(geom_neighbor.coords[0])) or pt.equals(Point(geom_neighbor.coords[-1])):
                        continue
                    # 把附近的结点捕捉过来，自身不变
                    if cal_dist(pt.coords[0], geom_neighbor.coords[0]) > cal_dist(pt.coords[0], geom_neighbor.coords[-1]):
                        # 终点近，用终点去捕捉
                        fix_st = True
                    else:
                        fix_st = False
                    geom_new = rubber_snap_within_distance(pt, geom_neighbor, fix_st, self.transform_distance)
                    pk_i = gdf[pk_col].iloc[i]
                    self.table.update_attr(pk_col, pk_i, {geom_col: geom_new}, auto_commit=False)
                modified = True
        else:  # 扩大搜索范围，查找附近是否有唯一线段（端点或中间点）
            gdf = self.table.buffer_intersect_by_pt(coord, self.extend_threshold, and_where=f"and {pk_col} != {pk_self}")
            if gdf.shape[0] == 2:  # 只有两条（即情况不复杂，再给个机会），进一步计算coord到各线段的距离
                gdf["dist"] = gdf[geom_col].apply(lambda x: x.distance(Point(coord)))
                gdf = gdf.sort_values(by="dist", ascending=True).reset_index(drop=True)
                if gdf["dist"].iloc[0] * self.extend_ratio < gdf["dist"].iloc[1]:  # coord到最近线段的距离小于第二近距离的1/10，取最近线段
                    gdf = gdf.iloc[0:1]
            if gdf.shape[0] == 1:  # 有唯一线段
                # 把自身【最近】的【端点】捕捉到附近的线段上，不管输入的coord是起点还是终点；优先捕捉到端点上
                geom_base = gdf[geom_col].iloc[0]
                geom_self_new = self.my_extend(geom_base, geom_self)
                self.table.update_attr(pk_col, pk_self, {geom_col: geom_self_new})
                modified = True
        return modified

    def my_extend(self, geom_base, geom_self):
        """
        把geom_self的端点捕捉到geom_base上。
        · 使用geom_self离geom_base最近的端点
        · 若geom_self的端点离geom_base的端点的距离符合要求，优先捕捉到geom_base的端点上
        · 若不符合要求，捕捉到离geom_self端点最近的点上。
        """
        st_self = geom_self.coords[0]
        ed_self = geom_self.coords[-1]
        st_dist = geom_base.distance(Point(st_self))
        ed_dist = geom_base.distance(Point(ed_self))
        reverse_flag = False
        if min(st_dist, ed_dist) > self.extend_threshold:
            raise ExtendException(min(st_dist, ed_dist), self.extend_threshold)
        if st_dist <= ed_dist:  # 起点更近，翻转，使其变为终点更近
            reverse_flag = True
            geom_self = reverse_linestring(geom_self)
            ed_self = geom_self.coords[-1]
        # 调整为ed_self与geom_base更近后，将ed_self捕捉到geom_base上
        pt_proj = self.get_pt_target(ed_self, geom_base)
        assert cal_dist(ed_self, pt_proj.coords[0]) < self.extend_threshold
        geom_self_new = rubber_snap_within_distance(pt_proj, geom_self, True, self.transform_distance)

        if reverse_flag:
            geom_self_new = reverse_linestring(geom_self_new)
        return geom_self_new

    def get_pt_target(self, coord, geom_base):
        # 确定捕捉点
        if self.end_priority:  # 如果端点也符合条件，优先捕捉到端点上
            if cal_dist(coord, geom_base.coords[0]) < self.extend_threshold:
                pt_target = Point(geom_base.coords[0])
            elif cal_dist(coord, geom_base.coords[-1]) < self.extend_threshold:
                pt_target = Point(geom_base.coords[-1])
            else:
                pt_target = nearest_points(Point(coord), geom_base)[1]
        else:
            pt_target = nearest_points(Point(coord), geom_base)[1]
        return pt_target


class FuseSnapTool(SnapTool):
    def snap_node(self, pk_self, coord, geom_self):
        pk_col = self.table.pk
        geom_col = self.table.geom
        gdf = self.table.buffer_nodes(Point(coord), self.tolerance, and_where=f"and {pk_col} != {pk_self}")
        modified = False
        if not gdf.empty:
            # 附近有结点
            with conn:
                for i in range(gdf.shape[0]):
                    pt = Point(coord)
                    geom_neighbor = gdf[geom_col].iloc[i]
                    if pt.equals(Point(geom_neighbor.coords[0])) or pt.equals(Point(geom_neighbor.coords[-1])):
                        continue
                    # 把附近的结点捕捉过来，自身不变
                    if cal_dist(pt.coords[0], geom_neighbor.coords[0]) > cal_dist(pt.coords[0], geom_neighbor.coords[-1]):
                        # 终点近，用终点去捕捉
                        fix_st = True
                    else:
                        fix_st = False
                    geom_new = rubber_snap_within_distance(pt, geom_neighbor, fix_st, self.transform_distance)
                    pk_i = gdf[pk_col].iloc[i]
                    self.table.update_attr(pk_col, pk_i, {geom_col: geom_new}, auto_commit=False)
                modified = True
        else:  # 扩大搜索范围，查找附近是否有唯一线段（端点或中间点）
            gdf = self.table.buffer_intersect_by_pt(coord, self.extend_threshold,
                                                    and_where=f"and {pk_col} != {pk_self}")
            if gdf.shape[0] == 2:  # 只有两条（即情况不复杂，再给个机会），进一步计算coord到各线段的距离
                gdf["dist"] = gdf[geom_col].apply(lambda x: x.distance(Point(coord)))
                gdf = gdf.sort_values(by="dist", ascending=True).reset_index(drop=True)
                if gdf["dist"].iloc[0] * self.extend_ratio < gdf["dist"].iloc[1]:  # coord到最近线段的距离小于第二近距离的1/10，取最近线段
                    gdf = gdf.iloc[0:1]
            if gdf.shape[0] == 1:  # 有唯一线段
                # 把【coord对应的端点】捕捉到附近的线段上；优先捕捉到端点上
                geom_base = gdf[geom_col].iloc[0]
                # 判断coord对应geom_self的起点还是终点
                if is_same_point(coord, geom_self.coords[-1]):  # coord是geom_self的终点，geom_self的起点不动，终点端捕捉
                    fix_st = True
                else:
                    fix_st = False
                # 确定捕捉点
                pt_target = self.get_pt_target(coord, geom_base)
                # 捕捉
                geom_self_new = rubber_snap_within_distance(pt_target, geom_self, fix_st, self.transform_distance)
                self.table.update_attr(pk_col, pk_self, {geom_col: geom_self_new})
                modified = True
        return modified
